import os
# 导入LangChain相关模块用于构建PDF问答系统
from langchain_community.document_loaders import PyPDFLoader      # PDF文档加载器
from langchain_community.vectorstores import FAISS               # 向量数据库
from langchain_classic.chains.retrieval_qa.base import RetrievalQA                         # 检索问答链
from langchain_text_splitters import RecursiveCharacterTextSplitter  # 文本分割器
# 导入OpenAI兼容的嵌入模型和聊天模型（用于DeepSeek API）
from langchain_openai import OpenAIEmbeddings, ChatOpenAI

# 设置 DeepSeek API 密钥 (从官网获取)
os.environ["DEEPSEEK_API_KEY"] = "sk-1ba4abc00cf84af689d207ae24299795"  # 替换为你的实际密钥

# 1. 加载 PDF 文档
def load_pdf_documents(file_path):
    """
    加载并处理PDF文档
    使用PyPDFLoader加载PDF文件，并通过RecursiveCharacterTextSplitter进行文本分割
    """
    # 创建PDF加载器实例
    loader = PyPDFLoader(file_path)
    # 加载PDF文档
    documents = loader.load()

    # 中文优化文本分割
    # 使用适合中文文档的分块策略，确保语义完整性
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,  # 适合中文文档的块大小
        chunk_overlap=200,  # 块之间重叠200个字符，确保上下文连贯性
        separators=["\n\n", "\n", "。", "！", "？", "．", "；", "，", " ", ""]  # 中文分段符
    )
    # 返回分割后的文档列表
    return text_splitter.split_documents(documents)


# 2. 创建向量数据库
def create_vector_store(docs):
    """
    创建文档向量存储
    使用FAISS向量数据库存储文档嵌入向量，便于后续相似性检索
    """
    # 使用 OpenAI 兼容的嵌入模型 (DeepSeek 暂未提供专用嵌入模型)
    embeddings = OpenAIEmbeddings(
        model="text-embedding-3-large",        # 使用高质量的嵌入模型
        openai_api_base="https://api.deepseek.com/v1",  # DeepSeek API 端点
        openai_api_key=os.environ["DEEPSEEK_API_KEY"]   # 使用环境变量中的API密钥
    )

    # 从文档创建FAISS向量数据库
    return FAISS.from_documents(docs, embeddings)


# 3. 创建 DeepSeek 问答链
def create_qa_chain(vector_store):
    """
    创建 DeepSeek 问答链
    使用DeepSeek API作为大语言模型，构建检索问答链
    """
    # 使用 ChatOpenAI 类来调用 DeepSeek API
    llm = ChatOpenAI(
        model="deepseek-chat",  # 使用 deepseek-chat 模型
        temperature=0.3,  # 控制创造性 (0-1)，较低值使输出更确定性
        max_tokens=2048,  # 最大输出长度
        base_url="https://api.deepseek.com/v1",  # API基础URL
        api_key=os.environ["DEEPSEEK_API_KEY"],  # API密钥
        timeout=30  # 增加超时时间，确保API调用完成
    )

    # 自定义提示模板 (中文优化)
    from langchain_core.prompts import PromptTemplate
    template = """
    你是一个专业的文档分析助手，请根据以下上下文回答问题：
    {context}

    问题：{question}

    要求：
    1. 答案必须基于提供的上下文
    2. 使用简洁的中文回答
    3. 如果上下文没有相关信息，回答"未在文档中找到相关信息"
    """
    # 创建提示模板实例
    prompt = PromptTemplate(template=template, input_variables=["context", "question"])

    # 创建检索问答链
    return RetrievalQA.from_chain_type(
        llm=llm,                                   # 使用的大型语言模型
        chain_type="stuff",                        # 链类型，将所有相关文档一次性传递给模型
        retriever=vector_store.as_retriever(search_kwargs={"k": 4}),  # 检索器，检索4个相关片段
        return_source_documents=True,              # 返回源文档
        chain_type_kwargs={"prompt": prompt}       # 使用自定义提示模板
    )


# 4. 主函数
def main():
    """
    主函数
    控制整个PDF问答流程：加载文档 -> 创建向量存储 -> 构建问答链 -> 执行问答
    """
    # 加载文档
    pdf_path = "../example.pdf"  # 替换为你的PDF文件路径
    docs = load_pdf_documents(pdf_path)
    print(f"✅ 已加载并分割文档，共 {len(docs)} 个文本块")

    # 创建向量存储
    vector_store = create_vector_store(docs)
    print("✅ 已创建向量数据库")

    # 创建问答链
    qa_chain = create_qa_chain(vector_store)
    print("✅ 已创建DeepSeek问答链\n")

    # 示例问题
    questions = [
        "本文档的主要研究内容是什么？",
        "给我罗列出营业相关数据？",
        "净利润和毛利分别是多少？"
    ]

    # 执行问答
    for question in questions:
        print(f"📝 问题: {question}")
        try:
            # 调用问答链获取答案
            result = qa_chain.invoke({"query": question})

            print(f"💡 回答: {result['result']}")
            print("🔍 来源:")
            # 显示前2个相关文档片段及其页码
            for i, doc in enumerate(result['source_documents'][:2]):  # 显示前2个来源
                print(f"  片段 {i + 1} (P{doc.metadata['page'] + 1}): {doc.page_content[:100]}...")
        except Exception as e:
            print(f"❌ 错误: {str(e)}")
        print("-" * 50)


# 程序入口点
if __name__ == "__main__":
    main()